import numpy as np
import matplotlib.pyplot as plt
import pandas as pd


def sep(label=''):
    """Utility function to print separator line."""
    print('-' * 32, label, '-' * 32, sep='')


# 1.	完成如下K-means算法底层操作（25分）
# (1)	数据初始化
# ①	读取drink.txt， 获取前两项特征（2分）
sep('①	读取drink.txt， 获取前两项特征')
df = pd.read_csv(r'../../../large_data/ML2/monthly_exam/drink.txt').iloc[:, :2]
m, n = df.shape
print(m, n)
print(df[:5])

# ②	实现数据归一化处理（底层）（2分）
sep('②	实现数据归一化处理（底层）')
xmax = df.max(axis=0)
xmin = df.min(axis=0)
df = (df - xmin) / (xmax - xmin)
print(df[:5])

# (2)	底层实现
# ①	随机初始化2个质心，并设置kmeans循环10次，K值为2（2分）
sep('①	随机初始化2个质心，并设置kmeans循环10次，K值为2')
iters = 10
k = 2
np.random.seed(666)
idx_rand = np.random.permutation(m)[:k]
centers = df.iloc[idx_rand].to_numpy()
print(f'初始化2个质心:{centers}')

for i in range(iters):
    # ②	计算各质心到样本之间的距离（2分）
    dis_mat = np.zeros([m, k])  # distance matrix
    for j in range(k):
        dis_mat[:, j] = ((df - centers[j]) ** 2).sum(axis=1) ** 0.5

    # ③	计算每个样本点的归属于哪一个簇（2分）
    cls_vector = dis_mat.argmin(axis=1)  # cluster vector

    # ④	更新质心位置（3分）
    centers_new = np.zeros([k, 2])
    for j in range(k):
        centers_new[j] = df.loc[cls_vector == j].mean(axis=0)

    # ⑤	如果当前质心和上一次更新位置不变，提前停止循环（3分）
    if np.allclose(centers, centers_new):
        print(f'提前停止循环在第{i + 1}次迭代。')
        break
    else:
        centers = centers_new
centers = centers_new

# ⑥	打印最后的质心坐标，及每个样本点的聚类归属（3分）
sep('⑥	打印最后的质心坐标，及每个样本点的聚类归属')
print(f'质心坐标:{centers}')
print(f'每个样本点的聚类归属:{cls_vector}')

# (3)	图像绘制
# ①	绘制样本点位置，每个簇使用一种颜色在图上显示（3分）
plt.scatter(df.iloc[:, 0], df.iloc[:, 1], c=cls_vector)

# ②	绘制出每个簇的质心位置（3分）
plt.scatter(centers[:, 0], centers[:, 1], marker='x', color='k')
for i in range(k):
    plt.annotate(f'Center#{i + 1}', xy=[centers[i, 0], centers[i, 1]])

# Finally show all plotting
plt.show()
